import numpy as np
import cv2

"""
code2:
    光流估计法，提取特征值并且检测运动距离
:param
    无
:return
    无
"""
# 加载视频
cap = cv2.VideoCapture()
cap.open(r'C:\OpenCV\data\vtest.avi')
if not cap.isOpened():
    print("can't open video \n")

# 方便调参，使用dict字典


# lucas kanade光流法参数设定
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
                 # winSize：每个金字塔级别的搜索窗口大小。
                 # maxLevel：从0开始的最大金字塔等级数字；如果设置为0，则不使用金字塔（单个级别）；如果设置为1，则使用两个级别。
                 # criteria：参数，指定迭代搜索算法的终止条件
                 # （在指定的最大迭代次数之后 criteria.maxCount 或当搜索窗口移动小于时 criteria.epsilon）。
# 角点检测参数设定
feature_params = dict(maxCorners=500,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)
                      # maxCorners : 设置最多返回的关键点数量。
                      # qualityLevel : 角点的品质因子，决定能不能成为角点。
                      # minDistance : 关键点之间的最少像素点，此邻域范围内如果存在更强角点，则删除此角点。
                      # blockSize : 计算一个角点是否为关键点时所取的区域大小。
# 读取第一帧视频图像
status, prev_img = cap.read()
if not status:
    print("can't read first video picture")
prev_gray = cv2.cvtColor(prev_img, cv2.COLOR_BGR2GRAY)
# 提取第一帧的特征点
p0 = cv2.goodFeaturesToTrack(prev_gray, mask=None, **feature_params)

while True:
    # 读取当前帧视频图像
    status, img = cap.read()
    if not status:
        print("can't read video \n")
        break
    cur_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 计算光流
    p1, st, err = cv2.calcOpticalFlowPyrLK(prev_gray, cur_gray, p0, None,
                                           **lk_params)
    # 选取好的跟踪点
    good_points = p1[st == 1]
    good_prv_points = p0[st == 1]
    # 在结果图像中叠加画出特征点和计算出来的光流向量
    res = img.copy()
    draw_line = (0, 0, 255)
    draw_circle = (0, 255, 0)
    # draw_color = np.random.randint(0, 255, (100,3))
    # zip将good, good_prv合成一个元祖
    for i, (cur, prev) in enumerate(zip(good_points,good_prv_points)):
        x0, y0 = cur.ravel()  # 将多维度转化为一维
        x1, y1 = prev.ravel()
        # 标注上特征点和运动轨迹
        # cv2.line(res, (x0,y0), (x1,y1), draw_color[i].tolist())
        # cv2.circle(res, (x0, y0), 3, draw_color[i].tolist())
        cv2.line(res, (x0,y0), (x1,y1), draw_line)
        cv2.circle(res, (x0, y0), 3, draw_circle)

    #更新上一帧
    prev_gray = cur_gray.copy()
    p0 = good_points.reshape(-1, 1, 2)

    # 显示计算结果图像
    cv2.imshow("check result pic", res)

    # 等待esc退出
    key = cv2.waitKey(0)
    if key == 27:
        break

cap.release()
cv2.destroyAllWindows()

